Remove getEnvEnablePDL in favor of enable_pdl parameter #1446

yongwww · 2025-08-10T17:45:11Z

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

I have installed pre-commit by running pip install pre-commit (or used your preferred method).
I have installed the hooks with pre-commit install.
I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

Tests have been added or updated as needed.
All tests are passing (unittest, etc.).

Reviewer Notes

gemini-code-assist

Summary of Changes

Hello @yongwww, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

I've refactored the mechanism for enabling Programmatic Dependent Launch (PDL) within the codebase. Previously, this feature was controlled by checking a global environment variable. Now, I've removed that environment variable check and instead introduced an explicit enable_pdl boolean parameter. This parameter is passed directly to all relevant CUDA kernel launch functions and associated utility calls across various components, including MoE, quantization, and GEMM operations. This change makes the control over PDL more explicit, localized, and easier to manage programmatically, improving clarity and reducing reliance on global state.

Highlights

Removal of Global Environment Variable Check: The global getEnvEnablePDL function and its declarations have been completely removed from the codebase.
Introduction of Explicit enable_pdl Parameter: A new enable_pdl boolean parameter has been added to numerous CUDA kernel launch functions and related helper functions across various modules, including Mixture-of-Experts (MoE) operations, quantization routines, and GEMM (General Matrix Multiply) implementations.
Struct Updates for PDL Control: The TmaWarpSpecializedGroupedGemmInput and TllmGenFmhaRunnerParams structs have been updated to include the enable_pdl flag, allowing this setting to be carried through data structures.
Propagation of PDL Control Through Call Stack: The enable_pdl parameter is now explicitly propagated through the entire call stack, from Python wrappers down to the underlying CUDA kernel launches, ensuring direct control over the Programmatic Dependent Launch feature.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request refactors the codebase to remove the dependency on the getEnvEnablePDL function, replacing it with an explicit enable_pdl parameter. This is a positive change that improves code clarity and testability by making dependencies explicit. The changes are extensive and mostly mechanical. I've identified a couple of critical issues where the new parameter was not propagated correctly through the call stack, which would result in compilation errors. Please see the detailed comments for suggestions on how to fix these issues.

gemini-code-assist · 2025-08-10T17:48:26Z

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -3005,7 +3007,7 @@ void CutlassMoeFCRunner<T, WeightType, OutputType, InputType, BackBoneType, Enab
        static_cast<UnfusedGemmOutputType const*>(gemm_output), fc2_fp8_quant, fc1_expert_biases,
        bias_is_broadcast, expert_first_token_offset, num_experts_per_node, inter_size,
        expanded_num_rows, fc1_activation_type, quant_params, use_per_expert_act_scale,
-        fc2_fp4_act_flat, stream);
+        fc2_fp4_act_flat, stream, enable_pdl);


The enable_pdl variable is used here, but it's not defined within the scope of gemm1. This will cause a compilation error. This issue also exists for gemm2 and its callees. The enable_pdl parameter needs to be plumbed through several functions.

Specifically, you need to:

Add bool enable_pdl to the function signatures of CutlassMoeFCRunner::gemm1 and CutlassMoeFCRunner::gemm2 in csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h and their implementations in this file.

Pass enable_pdl from runMoe to gemm1 and gemm2.

If BlockScaleFC1 and BlockScaleFC2 are called from gemm1 and gemm2 respectively, you'll need to add bool enable_pdl to their signatures and pass it down as well.

gemini-code-assist · 2025-08-10T17:48:26Z

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -4593,7 +4601,7 @@

  prepareRouting(num_tokens, workspace_ptr_char, stream);
  prepareQuantParams(num_tokens, workspace_ptr_char, stream);
-  prepareTmaWsInputs(num_tokens, workspace_ptr_char, expert_weights, stream);
+  prepareTmaWsInputs(num_tokens, workspace_ptr_char, expert_weights, stream, false);


The call to prepareRouting on line 4602 is missing the new enable_pdl argument. The function signature for prepareRouting was updated to include bool enable_pdl, but this call site was not updated, which will cause a compilation error. You should probably pass false here, similar to this call to prepareTmaWsInputs.

yzh119 · 2025-08-10T20:27:34Z

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

@@ -259,7 +259,8 @@ void buildMinLatencyActiveExpertMaps(int* num_active_experts_per_node,
                                     int const experts_per_token, int const start_expert,
                                     int const end_expert, int const num_experts_per_node,
                                     int const cluster_rank, int const cluster_size,
-                                     int const num_experts_smem, cudaStream_t const stream) {
+                                     int const num_experts_smem, cudaStream_t const stream,
+                                     bool enable_pdl = false) {


stream should be the last argument, and please do not set a default value for enable_pdl

csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh

zhyncs · 2025-08-12T21:11:51Z

@yongwww qq is this pr ready?

yzh119 · 2025-08-13T04:34:51Z

include/flashinfer/trtllm/fmha/fmhaRunnerParams.h

@@ -275,6 +275,8 @@ struct TllmGenFmhaRunnerParams {
  float mScaleSfO;
  // The cuda stream.
  cudaStream_t stream;
+  // Whether to enable PDL (Programmatic Dependent Launch).
+  bool enable_pdl = false;


Don't set default value here, it doesn't work, we will be memset the entire POD anyways in the constructors:
https://github.com/flashinfer-ai/flashinfer/pull/1446/files#diff-878c76ab5e42dac3beb66b2da51e8d72a9b994d44dee598f251fefdcad8b11dcR304

and for hopper/blackwell kernels, it should default to True, always.

nvpohanh · 2025-08-13T05:18:47Z

@yongwww @yzh119 I see that these only changed MoE-related kernels. Do we plan to make the same change for other kernels like Attn/Gemm/AR+Norm/activation/etc.?

yzh119 · 2025-08-13T06:09:35Z

@nvpohanh not only moe kernels, all dependency of getEnvEnablePDL (including trtllm-gen's batched gemm and attention) have been removed.

AR+Norm/activation

These operators do not rely on getEnvEnablePDL and we use enable_pdl when they were introduced to flashinfer from the beginning.

yzh119

CI is temporarily broken, run blackwell tests locally and all UT passed.

Thanks for your contribution @cyx-6 @yongwww !

IwakuraRein · 2025-08-13T23:52:04Z

@yzh119 @yongwww core.py:1442 is lacking the enable_pdl. I've created #1480 to fix it.

## 📌 Description Following #1446. ## 🔍 Related Issues  ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes  --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

## 📌 Description Fix the missing `enable_pdl` argument introduced in #1446 . ## 🔍 Related Issues  ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes  --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Avery Yingyi Huang <[email protected]>

Remove getEnvEnablePDL in favor of enable_pdl parameter

1250bf1

gemini-code-assist bot reviewed Aug 10, 2025

View reviewed changes

tqchen assigned cyx-6 Aug 10, 2025

yzh119 reviewed Aug 10, 2025

View reviewed changes

zhyncs added the priority: high label Aug 12, 2025

cyx-6 added 3 commits August 12, 2025 22:55

wip

43dd149

quant compilation passed

e32f724

fix moe

7ff9667

yzh119 marked this pull request as ready for review August 13, 2025 04:33

yzh119 reviewed Aug 13, 2025

View reviewed changes

cyx-6 added 2 commits August 13, 2025 04:36

fix lint

7328a03

upd

b74fbf3

fix

813c428

yzh119 approved these changes Aug 13, 2025

View reviewed changes

yzh119 merged commit 3628a54 into flashinfer-ai:main Aug 13, 2025
1 check passed

yongwww deleted the remove_pdl_env branch August 13, 2025 16:41

IwakuraRein mentioned this pull request Aug 13, 2025

fix missing enable_pdl argument in trtllm-gen fp4 moe #1480

Merged

5 tasks

yyihuang mentioned this pull request Aug 14, 2025

feat: add pdl for trtllm-gen attn #1484

Merged

5 tasks

Remove getEnvEnablePDL in favor of enable_pdl parameter #1446

Remove getEnvEnablePDL in favor of enable_pdl parameter #1446

Uh oh!

Conversation

yongwww commented Aug 10, 2025

📌 Description

🔍 Related Issues

🚀 Pull Request Checklist

✅ Pre-commit Checks

🧪 Tests

Reviewer Notes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Aug 10, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Aug 10, 2025

Choose a reason for hiding this comment

Uh oh!

yzh119 Aug 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhyncs commented Aug 12, 2025

Uh oh!

yzh119 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

cyx-6 Aug 13, 2025

Choose a reason for hiding this comment

Uh oh!

nvpohanh commented Aug 13, 2025

Uh oh!

yzh119 commented Aug 13, 2025

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

IwakuraRein commented Aug 13, 2025

Uh oh!

Uh oh!